#include "Aggregator.h"

#include <folly/logging/xlog.h>

namespace nexus {

Aggregator::Aggregator(Pipeline& pipeline) : pipeline_(pipeline) {}

void
Aggregator::run() {
  while (1) {
    std::shared_ptr<AtomicOverwrite> ao;
    auto ticket = pipeline_.blockingReadStage<2>(ao);
    auto out = std::make_shared<AtomicOverwrite>();
    out->device = ao->device;
    process(*ao, *out);
    pipeline_.blockingWriteStage<2>(ticket, std::move(out));
  }
}

void
Aggregator::process(AtomicOverwrite& ao, AtomicOverwrite& out) {
  std::map<Port*, std::vector<HeaderSpace>> grouped;
  for (auto& o : ao.overwrites) {
    grouped[o.second].push_back(std::move(o.first));
    // grouped[o.second] = grouped[o.second] | o.first;
  }

  // TODO: check if really need to group by pred, group by port seems ok for
  // computing AP
  for (auto& [k, v] : grouped) {
    auto res = std::reduce(
        std::execution::par_unseq,
        v.begin(),
        v.end(),
        HeaderSpace::hsEmpty(),
        [](auto x, auto y) { return x | y; });
    out.overwrites.push_back(std::make_pair(res, k));
    // out.overwrites.push_back(std::make_pair(HeaderSpace::orAll(v), k));
  }
}

} // namespace nexus